"""
    Thinning

Observation thinning to reduce correlated observation errors and computational cost.
Implements mesh-based thinning, priority-based selection, and superobbing.

# Motivation

Dense observations often have correlated errors, violating the independence
assumption in data assimilation. Thinning reduces:
1. Computational cost (fewer observations to process)
2. Correlation effects (spatial separation)
3. Imbalance in observation distribution

# Methods

## Mesh-Based Thinning
Divide domain into grid boxes, keep best observation per box.

## Priority-Based Selection
Among observations in same box, select by:
1. Quality (highest quality first)
2. Error (smallest error first)
3. Innovation (smallest departure from background)
4. Platform (priority ranking)

## Superobbing
Average multiple observations into single "super-observation"
with reduced error.

# Thinning Parameters

**Satellite radiances**: 50-100 km spacing
**Conventional obs**: 100-200 km spacing
**Radar**: 10-25 km spacing
**GPS-RO**: 200-300 km spacing
"""
module Thinning

using LinearAlgebra
using Statistics
using ...ObservationTypes

# Export types and functions
export ThinningConfig, ThinningResult, ThinningMethod
export perform_thinning, mesh_based_thinning, priority_based_selection
export create_superobs

@enum ThinningMethod begin
    MESH_BEST_QUALITY
    MESH_SMALLEST_ERROR
    MESH_SMALLEST_INNOVATION
    MESH_RANDOM
    PRIORITY_WEIGHTED
end

"""
    ThinningConfig{T<:AbstractFloat}

Configuration for observation thinning.
"""
struct ThinningConfig{T<:AbstractFloat}
    method::ThinningMethod
    mesh_size_km::T
    min_spacing_km::T
    max_obs_per_cell::Int
    platform_priorities::Dict{String,Int}
    quality_threshold::T
    create_superobs::Bool
    superobs_min_count::Int

    function ThinningConfig{T}(;
        method = MESH_BEST_QUALITY,
        mesh_size_km = T(100.0),
        min_spacing_km = T(50.0),
        max_obs_per_cell = 1,
        quality_threshold = T(0.5),
        create_superobs = false,
        superobs_min_count = 3
    ) where T
        # Default platform priorities (higher = better)
        priorities = Dict{String,Int}(
            "radiosonde" => 10,
            "aircraft" => 8,
            "surface" => 7,
            "ship" => 6,
            "buoy" => 6,
            "satellite" => 5,
            "gps_ro" => 7,
            "radar" => 4,
        )

        new{T}(method, mesh_size_km, min_spacing_km, max_obs_per_cell,
               priorities, quality_threshold, create_superobs, superobs_min_count)
    end
end

ThinningConfig() = ThinningConfig{Float64}()

"""
    ThinningResult{T<:AbstractFloat}

Result of observation thinning.
"""
struct ThinningResult{T<:AbstractFloat}
    thinned_observations::Vector{AbstractObservation{T}}
    kept_indices::Vector{Int}
    original_count::Int
    thinned_count::Int
    rejection_count::Int
    thinning_ratio::T
    statistics::Dict{String,Any}
end

"""
    perform_thinning(observations::Vector{<:AbstractObservation{T}},
                    config::ThinningConfig{T}) where T

Perform observation thinning.
"""
function perform_thinning(observations::Vector{<:AbstractObservation{T}},
                         config::ThinningConfig{T}) where T
    if isempty(observations)
        return ThinningResult{T}(
            AbstractObservation{T}[],
            Int[],
            0, 0, 0, T(0), Dict{String,Any}()
        )
    end

    original_count = length(observations)

    # Create spatial grid
    grid = create_thinning_grid(observations, config)

    # Select observations based on method
    kept_indices = if config.method == MESH_BEST_QUALITY
        select_best_quality(observations, grid, config)
    elseif config.method == MESH_SMALLEST_ERROR
        select_smallest_error(observations, grid, config)
    elseif config.method == MESH_SMALLEST_INNOVATION
        select_smallest_innovation(observations, grid, config)
    elseif config.method == PRIORITY_WEIGHTED
        priority_based_selection(observations, grid, config)
    else
        select_random(observations, grid, config)
    end

    thinned_obs = observations[kept_indices]
    thinned_count = length(thinned_obs)
    rejection_count = original_count - thinned_count

    stats = Dict{String,Any}(
        "original_count" => original_count,
        "thinned_count" => thinned_count,
        "rejection_count" => rejection_count,
        "thinning_ratio" => thinned_count / max(1, original_count),
        "grid_cells" => length(grid),
        "avg_obs_per_cell" => thinned_count / max(1, length(grid))
    )

    return ThinningResult{T}(
        thinned_obs,
        kept_indices,
        original_count,
        thinned_count,
        rejection_count,
        T(thinned_count / max(1, original_count)),
        stats
    )
end

"""
    create_thinning_grid(observations, config)

Create spatial grid for thinning.
"""
function create_thinning_grid(observations::Vector{<:AbstractObservation{T}},
                             config::ThinningConfig{T}) where T
    if isempty(observations)
        return Dict{Tuple{Int,Int},Vector{Int}}()
    end

    # Find domain bounds
    lats = [obs.location.latitude for obs in observations]
    lons = [obs.location.longitude for obs in observations]

    lat_min, lat_max = extrema(lats)
    lon_min, lon_max = extrema(lons)

    # Grid spacing in degrees (approximate)
    deg_per_km = 1.0 / 111.0  # Approximate at equator
    grid_spacing = config.mesh_size_km * deg_per_km

    # Create grid cells
    grid = Dict{Tuple{Int,Int},Vector{Int}}()

    for (i, obs) in enumerate(observations)
        # Compute grid cell indices
        lat_idx = floor(Int, (obs.location.latitude - lat_min) / grid_spacing)
        lon_idx = floor(Int, (obs.location.longitude - lon_min) / grid_spacing)

        cell_key = (lat_idx, lon_idx)

        if haskey(grid, cell_key)
            push!(grid[cell_key], i)
        else
            grid[cell_key] = [i]
        end
    end

    return grid
end

"""
    select_best_quality(observations, grid, config)

Select observation with best quality in each grid cell.
"""
function select_best_quality(observations::Vector{<:AbstractObservation{T}},
                            grid::Dict{Tuple{Int,Int},Vector{Int}},
                            config::ThinningConfig{T}) where T
    kept = Int[]

    for (cell, indices) in grid
        if isempty(indices)
            continue
        end

        # Get quality for each observation in cell
        qualities = [get_observation_quality(observations[i]) for i in indices]

        # Select best quality
        best_idx = argmax(qualities)

        # Check if meets quality threshold
        if qualities[best_idx] >= config.quality_threshold
            push!(kept, indices[best_idx])
        end
    end

    return sort(kept)
end

"""
    select_smallest_error(observations, grid, config)

Select observation with smallest error in each grid cell.
"""
function select_smallest_error(observations::Vector{<:AbstractObservation{T}},
                              grid::Dict{Tuple{Int,Int},Vector{Int}},
                              config::ThinningConfig{T}) where T
    kept = Int[]

    for (cell, indices) in grid
        if isempty(indices)
            continue
        end

        # Get errors
        errors = [observations[i].observation_error.total_error for i in indices]

        # Select smallest error
        best_idx = argmin(errors)
        push!(kept, indices[best_idx])
    end

    return sort(kept)
end

"""
    select_smallest_innovation(observations, grid, config)

Select observation with smallest background departure.
"""
function select_smallest_innovation(observations::Vector{<:AbstractObservation{T}},
                                   grid::Dict{Tuple{Int,Int},Vector{Int}},
                                   config::ThinningConfig{T}) where T
    kept = Int[]

    for (cell, indices) in grid
        if isempty(indices)
            continue
        end

        # Get innovations (obs - background)
        innovations = [abs(observations[i].observed_value - observations[i].background_value)
                      for i in indices]

        # Select smallest innovation
        best_idx = argmin(innovations)
        push!(kept, indices[best_idx])
    end

    return sort(kept)
end

"""
    select_random(observations, grid, config)

Randomly select one observation per grid cell.
"""
function select_random(observations::Vector{<:AbstractObservation{T}},
                      grid::Dict{Tuple{Int,Int},Vector{Int}},
                      config::ThinningConfig{T}) where T
    kept = Int[]

    for (cell, indices) in grid
        if isempty(indices)
            continue
        end

        # Random selection
        selected = rand(indices)
        push!(kept, selected)
    end

    return sort(kept)
end

"""
    priority_based_selection(observations, grid, config)

Select based on combined priority score.
"""
function priority_based_selection(observations::Vector{<:AbstractObservation{T}},
                                 grid::Dict{Tuple{Int,Int},Vector{Int}},
                                 config::ThinningConfig{T}) where T
    kept = Int[]

    for (cell, indices) in grid
        if isempty(indices)
            continue
        end

        # Compute priority scores
        scores = [compute_priority_score(observations[i], config) for i in indices]

        # Select highest priority
        best_idx = argmax(scores)
        push!(kept, indices[best_idx])
    end

    return sort(kept)
end

"""
    compute_priority_score(obs, config)

Compute priority score combining quality, error, and platform.
"""
function compute_priority_score(obs::AbstractObservation{T}, config::ThinningConfig{T}) where T
    # Quality component (0-1)
    quality_score = get_observation_quality(obs)

    # Error component (inverse of normalized error)
    error_score = T(1.0) / max(T(0.1), obs.observation_error.total_error)

    # Platform priority (normalized to 0-1)
    platform_type = get_platform_type(obs)
    platform_score = get(config.platform_priorities, platform_type, 5) / T(10.0)

    # Combined score (weighted sum)
    score = T(0.4) * quality_score + T(0.3) * error_score + T(0.3) * platform_score

    return score
end

"""
    get_observation_quality(obs)

Extract quality metric from observation.
"""
function get_observation_quality(obs::AbstractObservation{T}) where T
    # Check QC flags
    if obs.quality_control.analysis_usage_flag < 0
        return T(0.0)  # Rejected
    elseif obs.quality_control.analysis_usage_flag == 0
        return T(0.5)  # Monitoring
    end

    # Quality based on QC flags
    quality = T(1.0)

    if obs.quality_control.gross_error_flag > 0
        quality *= T(0.8)
    end

    if obs.quality_control.background_check_flag > 0
        quality *= T(0.9)
    end

    if obs.quality_control.buddy_check_flag > 0
        quality *= T(0.9)
    end

    return quality
end

"""
    get_platform_type(obs)

Get platform type string.
"""
function get_platform_type(obs::SurfaceObservation)::String
    return "surface"
end

function get_platform_type(obs::RadiosondeObservation)::String
    return "radiosonde"
end

function get_platform_type(obs::AircraftObservation)::String
    return "aircraft"
end

function get_platform_type(obs::RadianceObservation)::String
    return "satellite"
end

function get_platform_type(obs::GPSROObservation)::String
    return "gps_ro"
end

function get_platform_type(obs::RadarObservation)::String
    return "radar"
end

function get_platform_type(obs::AbstractObservation)::String
    return "unknown"
end

end # module Thinning
